import java.io.*;
import java.util.*;

import org.jdom.*;
import org.jdom.filter.*;
import org.jdom.input.*;
import org.jdom.output.*;


public class ACEDocument {

  String m_strId;
  String m_strDir;

  Vector m_vecEntities;
  Vector m_vecRelations;

  String m_strText;

  public ACEDocument()
  {
    m_vecEntities = new Vector();
    m_vecRelations = new Vector();

    m_strText = new String();
  }
  
  public void load(String strFile)
  {

    int nPos = strFile.lastIndexOf("/");
    m_strDir = strFile.substring(0, nPos);
    m_strId = strFile.substring(nPos);

    loadXML("/home/pjbarrio/Dataset/ACE2003/ace_phase2/data/train/" + m_strId + ".sgm.tmx.rdc.xml");
    
    loadSGM(m_strDir + "/" + m_strId + ".sgm");
  }


  public void removeRepetitions()
  {
    for (int i = 0; i < m_vecEntities.size(); i++) {
      ACEEntity e1 = (ACEEntity) m_vecEntities.get(i);
      ListIterator it1 = e1.m_lstMentions.listIterator();
      while (it1.hasNext()) {
	ACEEMention em1 = (ACEEMention) it1.next();
	for (int j = i + 1; j < m_vecEntities.size(); j++) {
	  ACEEntity e2 = (ACEEntity) m_vecEntities.get(j);
	  ListIterator it2 = e2.m_lstMentions.listIterator();
	  while (it2.hasNext()) {
	    ACEEMention em2 = (ACEEMention) it2.next();
	    if (em1.compareTo(em2) == 0) {
	      it2.remove();
	      // replace em2 with em1 in all relations
	      for (int k = 0; k < m_vecRelations.size(); k++) {
		ACERelation r = (ACERelation) m_vecRelations.get(k);
		ListIterator itr = r.m_lstMentions.listIterator();
		while (itr.hasNext()) {
		  ACERMention rm = (ACERMention) itr.next();
		  if (rm.m_arg1 == em2)
		    rm.m_arg1 = em1;
		  if (rm.m_arg2 == em2)
		    rm.m_arg2 = em1;
		}
	      }
	    }
	  }
	}
      }
    }
  }


  public void loadXML(String strXMLFile)
  {
    try {
      SAXBuilder builder = new SAXBuilder();
      builder.setExpandEntities(true);
      Document doc = builder.build(strXMLFile);

      Element rootElement = doc.getRootElement();
      TreeMap mapEntities = new TreeMap();
      TreeMap mapEMentions = new TreeMap();
      process(rootElement, mapEntities, mapEMentions);
    }
    catch (JDOMException e) {
      e.printStackTrace();
    } 
    catch (IOException e) {
      e.printStackTrace();
    }   
  }


  public void loadSGM(String strSGMFile)
  {
    try {
      m_strText = "";
      
      BufferedReader br = new BufferedReader(new FileReader(strSGMFile));
      String strLine = br.readLine();
      while (strLine != null) {
	int i = 0;
	while (i < strLine.length()) {
	  char c = strLine.charAt(i);
	  if (c == '<')
	    i = strLine.indexOf('>', i);
	  else
	    m_strText += c;
	  i++;
	}
	m_strText += '\n';
	strLine = br.readLine();
      }
      br.close();
    }
    catch (IOException e) {
      e.printStackTrace();
    }   
  }


  public void process(Element elem, 
		      TreeMap mapEntities, TreeMap mapEMentions)
  {
    String strName = elem.getName();
    if (strName.equals("document"))
      m_strId = elem.getAttribute("DOCID").getValue();
    
    if (strName.equals("entity"))
      processEntity(elem, mapEntities, mapEMentions);
    else if (strName.equals("relation"))
      processRelation(elem, mapEntities, mapEMentions);
    else {
      List lstChildren = elem.getChildren();
      ListIterator it = lstChildren.listIterator();
      while (it.hasNext()) {
	Element el = (Element) it.next();
	process(el, mapEntities, mapEMentions);
      }
    }
  }


  public void processEntity(Element elem, 
			    TreeMap mapEntities, TreeMap mapEMentions)
  {
    // process entity
    ACEEntity entity = new ACEEntity();
    entity.m_strID = elem.getAttribute("ID").getValue();
    m_vecEntities.add(entity);
    mapEntities.put(entity.m_strID, entity);

    // process entity mentions
    List lstChildren = elem.getChildren();
    ListIterator it = lstChildren.listIterator();
    while (it.hasNext()) {
      Element el = (Element) it.next();
      String strName = el.getName();
      if (strName.equals("entity_type")) {
	entity.m_strType = el.getTextTrim();
	entity.m_bGeneric = el.getAttribute("GENERIC").equals("TRUE");
      }
      else if (strName.equals("entity_mention")) {
	ACEEMention emention = new ACEEMention();
	emention.m_strID = el.getAttribute("ID").getValue();
	emention.m_strType = el.getAttribute("TYPE").getValue();
	setAttributes(emention, el);
	mapEMentions.put(emention.m_strID, emention);
	entity.addMention(emention);
      }
    }
  }


  protected void setAttributes(ACEEMention em, Element el)
  {
    List lst1 = el.getChildren();
    ListIterator it1 = lst1.listIterator();
    while (it1.hasNext()) {
      Element el1 = (Element) it1.next();
      if (el1.getName().equals("extent")) {
	List lst2 = el1.getChildren();
	ListIterator it2 = lst2.listIterator();
	while (it2.hasNext()) {
	  Element el2 = (Element) it2.next();
	  if (el2.getName().equals("charseq")) {
	    em.m_nBeg = Integer.valueOf(el2.getChildText("start")).intValue();
	    em.m_nLen = Integer.valueOf(el2.getChildText("end")).intValue() - 
	      em.m_nBeg + 1;
	  }
	}
      }
      else if (el1.getName().equals("head")) {
	List lst2 = el1.getChildren();
	ListIterator it2 = lst2.listIterator();
	while (it2.hasNext()) {
	  Element el2 = (Element) it2.next();
	  if (el2.getName().equals("charseq")) {
	    em.m_nHBeg = Integer.valueOf(el2.getChildText("start")).intValue();
	    em.m_nHLen = Integer.valueOf(el2.getChildText("end")).intValue() - 
	      em.m_nHBeg + 1;
	  }
	}
      }
    }
  }


  public void processRelation(Element elem, 
			      TreeMap mapEntities, TreeMap mapEMentions)
  {
    // process relation
    ACERelation relation = new ACERelation();
    Attribute attrID = elem.getAttribute("ID");
    relation.m_strID = attrID.getValue();
    relation.m_strType = elem.getAttribute("TYPE").getValue();
    relation.m_strSType = elem.getAttribute("SUBTYPE").getValue();
    String strClass = elem.getAttribute("CLASS").getValue();
    relation.m_bExplicit = strClass.equals("EXPLICIT");
    m_vecRelations.add(relation);

    // process relation arguments & mentions
    List lstChildren = elem.getChildren();
    ListIterator it = lstChildren.listIterator();
    while (it.hasNext()) {
      Element el = (Element) it.next();
      String strName = el.getName();
      if (strName.equals("rel_entity_arg")) { // process arguments
	String strID = el.getAttribute("ENTITYID").getValue();
	ACEEntity entArg = (ACEEntity) mapEntities.get(strID);
	try {
	  int nArg = el.getAttribute("ARGNUM").getIntValue();
	  if (nArg == 1)
	    relation.m_arg1 = entArg;
	  else if (nArg == 2)
	    relation.m_arg2 = entArg;
	  else
	    assert false;
	}
	catch (DataConversionException e) {
	  e.printStackTrace();
	}
      }
      else if (strName.equals("relation_mentions")) { // process mentions
	List lstMentions = el.getChildren();
	ListIterator lit = lstMentions.listIterator();
	while (lit.hasNext()) {
	  Element elMention = (Element) lit.next();
	  if (elMention.getName().equals("relation_mention")) {
	    ACERMention rmention = new ACERMention();
	    rmention.m_strID = elMention.getAttribute("ID").getValue();
	    ListIterator argit = elMention.getChildren().listIterator();
	    while (argit.hasNext()) {
	      Element elArg = (Element) argit.next();
	      if (elArg.getName().equals("rel_mention_arg")) {
		ACEEMention emArg = (ACEEMention) 
		  mapEMentions.get(elArg.getAttribute("MENTIONID").getValue());
		try {
		  int nArg = elArg.getAttribute("ARGNUM").getIntValue();
		  if (nArg == 1)
		    rmention.m_arg1 = emArg;
		  else if (nArg == 2)
		    rmention.m_arg2 = emArg;
		  else
		    assert false;
		}
		catch (DataConversionException e) {
		  e.printStackTrace();
		}
	      }
	    }

	    relation.addMention(rmention);
	  }
	}
      }
    }
  }


  public TreeMap computeDegrees()
  {
    TreeMap mapER = getEntityRelations();
    TreeMap mapERels = new TreeMap();

    Set entries = mapER.entrySet();
    Iterator it = entries.iterator();
    while (it.hasNext()) {
      Map.Entry entry = (Map.Entry) it.next();
      ACEEMention em = (ACEEMention) entry.getKey();
      mapERels.put(em.m_strID,
		   getRelationClasses(em, (LinkedList) entry.getValue()));
    }

    return mapERels;
  }

  
  protected LinkedList getRelationClasses(ACEEMention em,
					  LinkedList lstRelations)
  {
    // create ordered list of relations
    LinkedList lstRel = new LinkedList();
    ListIterator it1 = lstRelations.listIterator();
    while (it1.hasNext()) {
      ACERMention rm1 = (ACERMention) it1.next();
      ACEEMention em1 = getOtherEMention(rm1, em);
      ListIterator it2 = lstRel.listIterator();
      boolean bFound = false;
      while (it2.hasNext() && !bFound) {
	ACERMention rm2 = (ACERMention) it2.next();
	ACEEMention em2 = getOtherEMention(rm2, em);
	if (em1.compareTo(em2) < 0) {
	  it2.previous();
	  bFound = true;
	}
      }
      it2.add(rm1);
    }

    // create relation classes
    LinkedList lstClasses = new LinkedList();

    LinkedList lstLast = new LinkedList();
    ACEEMention emLast = null;

    ListIterator it = lstRel.listIterator();
    while (it.hasNext()) {
      ACERMention rmCrt = (ACERMention) it.next();
      ACEEMention emCrt = getOtherEMention(rmCrt, em);
      if (lstLast.size() == 0 || betweenComma(emLast, emCrt)) {
	lstLast.add(rmCrt);
	emLast = emCrt;
      }
      else if (betweenAnd(emLast, emCrt)) {
	lstLast.add(rmCrt);
	lstClasses.add(lstLast);
	lstLast = new LinkedList();
	emLast = null;
      }
      else {
	ListIterator lit = lstLast.listIterator();
	while (lit.hasNext()) {
	  LinkedList lstSingle = new LinkedList();
	  lstSingle.add(lit.next());
	  lstClasses.add(lstSingle);
	}
	lstLast = new LinkedList();
	lstLast.add(emCrt);
	emLast = emCrt;
      }
    }
    
    // process last relations
    ListIterator lit = lstLast.listIterator();
    while (lit.hasNext()) {
      LinkedList lstSingle = new LinkedList();
      lstSingle.add(lit.next());
      lstClasses.add(lstSingle);
    }

    return lstClasses;
  }


  protected ACEEMention getOtherEMention(ACERMention rm, ACEEMention em)
  {
    if (rm.m_arg1 == em)
      return rm.m_arg2;
    else
      return rm.m_arg1;
  }


  protected boolean betweenComma(ACEEMention em1, ACEEMention em2)
  {
    int nStart = em1.m_nBeg + em1.m_nLen;
    int nEnd = em2.m_nBeg;
    
    if (nStart > nEnd)
      return false;

    String strBetween = m_strText.substring(nStart, nEnd).trim();

    return strBetween.equals(",") || 
      strBetween.equals(";") ||
      strBetween.equals(".,") ||
      strBetween.equals(".;");
  }


  protected boolean betweenAnd(ACEEMention em1, ACEEMention em2)
  {
    int nStart = em1.m_nBeg + em1.m_nLen;
    int nEnd = em2.m_nBeg;

    if (nStart > nEnd)
      return false;

    String strBetween = m_strText.substring(nStart, nEnd).trim();

    if (strBetween.equals("and"))
      return true;

    if (strBetween.length() == 0 || 
	(strBetween.charAt(0) != ',' && 
	 strBetween.charAt(0) != ';' &&
	 strBetween.charAt(0) != '.'))
      return false;

    String strAnd = strBetween.substring(1).trim();
    return strAnd.equals("and");
  }


  protected TreeMap getEntityRelations()
  {
    TreeMap mapEM = new TreeMap();
    for (int i = 0; i < m_vecRelations.size(); i++) {
      ACERelation relation = (ACERelation) m_vecRelations.get(i);
      if (relation.isExplicit()) {
	ListIterator it = relation.m_lstMentions.listIterator();
	while (it.hasNext()) {
	  ACERMention rm = (ACERMention) it.next();
	  addRelation(mapEM, rm.m_arg1, rm);
	  addRelation(mapEM, rm.m_arg2, rm);
	}
      }
    }

    return mapEM;
  }


  protected void addRelation(TreeMap mapEM, ACEEMention em1, ACERMention rm)
  {
    LinkedList lstRM = (LinkedList) mapEM.get(em1);
    if (lstRM == null) {
      lstRM = new LinkedList();
      mapEM.put(em1, lstRM);
    }

    lstRM.add(rm);
  }


  public boolean overlap()
  {
    boolean bOverlap = false;

    for (int i = 0; i < m_vecRelations.size(); i++) {
      ACERelation r1 = (ACERelation) m_vecRelations.get(i);
      if (r1.isExplicit()) {
	ListIterator it1 = r1.m_lstMentions.listIterator();
	while (it1.hasNext()) {
	  ACERMention rm1 = (ACERMention) it1.next();
	  for (int j = i; j < m_vecRelations.size(); j++) {
	    ACERelation r2 = (ACERelation) m_vecRelations.get(j);
	    if (r2.isExplicit()) {
	      ListIterator it2 = r2.m_lstMentions.listIterator();
	      while (it2.hasNext()) {
		ACERMention rm2 = (ACERMention) it2.next();
		if (rm1.overlap(rm2)) {
		  System.out.println("  Overlap between " + rm1.m_strID +
				     " and " + rm2.m_strID);
		  bOverlap = true;
		}
	      }
	    }
	  }
	}
      }
    }

    return bOverlap;
  }


  public void include()
  {
    for (int i = 0; i < m_vecRelations.size(); i++) {
      ACERelation r1 = (ACERelation) m_vecRelations.get(i);
      if (r1.isExplicit()) {
	ListIterator it1 = r1.m_lstMentions.listIterator();
	while (it1.hasNext()) {
	  ACERMention rm1 = (ACERMention) it1.next();
	  for (int j = i; j < m_vecRelations.size(); j++) {
	    ACERelation r2 = (ACERelation) m_vecRelations.get(j);
	    if (r2.isExplicit()) {
	      ListIterator it2 = r2.m_lstMentions.listIterator();
	      while (it2.hasNext()) {
		ACERMention rm2 = (ACERMention) it2.next();
		if (rm1.include(rm2)) 
		  rm1.m_nInc++;
		else if (rm2.include(rm1))
		  rm2.m_nInc++;
	      }
	    }
	  }
	}
      }
    }
  }


  public static void main_include(String[] args) 
  {
    if (args.length != 2) {
      System.out.println("Usage: java ACEDocument <corpus dir> <corpus file>!"); 
      return;
    }

    try {
      BufferedReader br = new BufferedReader(new FileReader(args[1]));
      String strLine = br.readLine();
      while (strLine != null) {
	System.out.println("Document " + strLine + ":");

	String strFile = args[0] + "/" + strLine.trim();
	ACEDocument aceDoc = new ACEDocument();
	aceDoc.load(strFile);

	aceDoc.include();
	// display includes
	for (int i = 0; i < aceDoc.m_vecRelations.size(); i++) {
	  ACERelation r1 = (ACERelation) aceDoc.m_vecRelations.get(i);
	  if (r1.isExplicit()) {
	    ListIterator it1 = r1.m_lstMentions.listIterator();
	    while (it1.hasNext()) {
	      ACERMention rm1 = (ACERMention) it1.next();
	      System.out.println("Relation " + rm1.m_strID + 
				 " includes: " + rm1.m_nInc);
	    }
	  }
	}

	/*
	// display degrees
	TreeMap mapEM = aceDoc.computeDegrees();
	Set entries = mapEM.entrySet();
	Iterator it = entries.iterator();
	while (it.hasNext()) {
	  Map.Entry entry = (Map.Entry) it.next();
	  LinkedList lstRels = (LinkedList) entry.getValue();
	  System.out.println("    Entity mention " + entry.getKey() + 
			     " has degree: " + lstRels.size());
	}
	*/

	strLine = br.readLine();
      }
    }
    catch (IOException e) {
      e.printStackTrace();
    }
  }
  

  public static void main_txt(String[] args) 
  {
    if (args.length != 2) {
      System.out.println("Usage: java ACEDocument <corpus dir> <corpus file>!"); 
      return;
    }

    try {
      BufferedReader br = new BufferedReader(new FileReader(args[1]));
      String strLine = br.readLine();
      while (strLine != null) {
	System.out.println("Document " + strLine + ":");

	String strFile = args[0] + "/" + strLine.trim();
	ACEDocument aceDoc = new ACEDocument();
	aceDoc.loadSGM(strFile + ".sgm");
	
	String strTXTFile = strFile + ".txt";
	PrintWriter pw = new PrintWriter(new FileWriter(strTXTFile));
	pw.print(aceDoc.m_strText);
	pw.close();

	strLine = br.readLine();
      }
    }
    catch (IOException e) {
      e.printStackTrace();
    }
  }
  

  public static void main(String[] args) 
  {
    if (args.length != 2) {
      System.out.println("Usage: java ACEDocument <corpus dir> <corpus list!"); 
      return;
    }

    int nCount = 0;
    int nExplicit = 0;
    try {
      BufferedReader br = new BufferedReader(new FileReader(args[1]));
      String strLine = br.readLine();
      while (strLine != null) {
	System.out.println("Document " + strLine + ":");

	String strFile = args[0] + "/" + strLine.trim();
	ACEDocument aceDoc = new ACEDocument();
	aceDoc.load(strFile);
	
	for (int i = 0; i < aceDoc.m_vecRelations.size(); i++) {
	  ACERelation aceRel = (ACERelation) aceDoc.m_vecRelations.get(i);
	  if (aceRel.m_bExplicit)
	    nExplicit += aceRel.m_lstMentions.size();
	  nCount += aceRel.m_lstMentions.size();
	}

	strLine = br.readLine();
      }
    }
    catch (IOException e) {
      e.printStackTrace();
    }

    System.out.println("Found " + nCount + " relation mentions!");
    System.out.println("Found " + nExplicit + " explicit mentions!");
  }
  

}
